#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <winsock2.h>
#include <ws2tcpip.h>

#pragma comment(lib, "ws2_32.lib")

#define MAX_CLIENTS 10
#define BUFFER_SIZE 2048

// 客户端信息结构体
typedef struct {
    SOCKET socket;
    char nickname[32];
    int is_connected;
} Client;

Client clients[MAX_CLIENTS];
int client_count = 0;
SOCKET server_socket;
fd_set master_set, read_set;

// 初始化Winsock库
int init_winsock() {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        printf("WSAStartup失败: %d\n", WSAGetLastError());
        return 0;
    }
    return 1;
}

// 创建并配置服务器套接字
int create_server_socket(int port) {
    server_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (server_socket == INVALID_SOCKET) {
        printf("创建套接字失败: %d\n", WSAGetLastError());
        return 0;
    }

    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);

    if (bind(server_socket, (struct sockaddr*)&server_addr, sizeof(server_addr)) == SOCKET_ERROR) {
        printf("绑定套接字失败: %d\n", WSAGetLastError());
        closesocket(server_socket);
        return 0;
    }

    if (listen(server_socket, 5) == SOCKET_ERROR) {
        printf("监听失败: %d\n", WSAGetLastError());
        closesocket(server_socket);
        return 0;
    }

    printf("服务器启动成功，监听端口 %d\n", port);
    return 1;
}

// 初始化客户端数组
void init_clients() {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        clients[i].socket = INVALID_SOCKET;
        clients[i].is_connected = 0;
        memset(clients[i].nickname, 0, sizeof(clients[i].nickname));
    }
}

// 添加新客户端
void add_client(SOCKET client_socket) {
    if (client_count >= MAX_CLIENTS) {
        printf("客户端数量已达上限，拒绝连接\n");
        closesocket(client_socket);
        return;
    }

    // 查找空位置
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (!clients[i].is_connected) {
            clients[i].socket = client_socket;
            clients[i].is_connected = 1;
            client_count++;
            
            // 接收客户端昵称
            char nickname[32];
            int recv_len = recv(client_socket, nickname, sizeof(nickname) - 1, 0);
            if (recv_len > 0) {
                nickname[recv_len] = '\0';
                strncpy(clients[i].nickname, nickname, sizeof(clients[i].nickname) - 1);
                printf("客户端 '%s' 已连接\n", clients[i].nickname);
            } else {
                strncpy(clients[i].nickname, "Anonymous", sizeof(clients[i].nickname) - 1);
                printf("客户端已连接 (匿名)\n");
            }
            
            // 通知其他客户端有新用户加入
            char welcome_msg[BUFFER_SIZE];
            sprintf(welcome_msg, "系统: 用户 '%s' 加入聊天室\n", clients[i].nickname);
            for (int j = 0; j < MAX_CLIENTS; j++) {
                if (clients[j].is_connected && j != i) {
                    send(clients[j].socket, welcome_msg, strlen(welcome_msg), 0);
                }
            }
            
            // 向新客户端发送欢迎消息
            sprintf(welcome_msg, "欢迎来到聊天室, %s!\n", clients[i].nickname);
            send(client_socket, welcome_msg, strlen(welcome_msg), 0);
            
            break;
        }
    }
}

// 移除客户端
void remove_client(int index) {
    if (clients[index].is_connected) {
        printf("客户端 '%s' 已断开连接\n", clients[index].nickname);
        
        // 通知其他客户端用户离开
        char leave_msg[BUFFER_SIZE];
        sprintf(leave_msg, "系统: 用户 '%s' 离开聊天室\n", clients[index].nickname);
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (clients[i].is_connected && i != index) {
                send(clients[i].socket, leave_msg, strlen(leave_msg), 0);
            }
        }
        
        closesocket(clients[index].socket);
        clients[index].socket = INVALID_SOCKET;
        clients[index].is_connected = 0;
        memset(clients[index].nickname, 0, sizeof(clients[index].nickname));
        client_count--;
    }
}

// 广播消息给所有客户端
void broadcast_message(int sender_index, char* message) {
    char broadcast_buffer[BUFFER_SIZE];
    sprintf(broadcast_buffer, "%s: %s", clients[sender_index].nickname, message);
    
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].is_connected && i != sender_index) {
            if (send(clients[i].socket, broadcast_buffer, strlen(broadcast_buffer), 0) == SOCKET_ERROR) {
                printf("发送消息给客户端 %d 失败\n", i);
            }
        }
    }
}

// 主循环处理客户端连接和消息
void run_server() {
    // 初始化文件描述符集合
    FD_ZERO(&master_set);
    FD_SET(server_socket, &master_set);
    
    int max_fd = server_socket;
    
    while (1) {
        // 复制主集合到读集合
        read_set = master_set;
        
        // 等待活动
        int activity = select(max_fd + 1, &read_set, NULL, NULL, NULL);
        if (activity == SOCKET_ERROR) {
            printf("select错误: %d\n", WSAGetLastError());
            break;
        }
        
        // 检查服务器套接字是否有新连接
        if (FD_ISSET(server_socket, &read_set)) {
            struct sockaddr_in client_addr;
            int addr_len = sizeof(client_addr);
            SOCKET new_client = accept(server_socket, (struct sockaddr*)&client_addr, &addr_len);
            
            if (new_client == INVALID_SOCKET) {
                printf("接受连接失败: %d\n", WSAGetLastError());
            } else {
                // 将新客户端添加到文件描述符集合
                FD_SET(new_client, &master_set);
                if (new_client > max_fd) {
                    max_fd = new_client;
                }
                
                add_client(new_client);
            }
            
            // 从读集合中移除服务器套接字
            FD_CLR(server_socket, &read_set);
        }
        
        // 检查每个客户端是否有消息
        for (int i = 0; i < MAX_CLIENTS; i++) {
            if (clients[i].is_connected && FD_ISSET(clients[i].socket, &read_set)) {
                char buffer[BUFFER_SIZE];
                int recv_len = recv(clients[i].socket, buffer, sizeof(buffer) - 1, 0);
                
                if (recv_len <= 0) {
                    // 客户端断开连接
                    FD_CLR(clients[i].socket, &master_set);
                    remove_client(i);
                } else {
                    // 处理接收到的消息
                    buffer[recv_len] = '\0';
                    printf("收到 %s 的消息: %s\n", clients[i].nickname, buffer);
                    broadcast_message(i, buffer);
                }
            }
        }
    }
}

// 清理资源
void cleanup() {
    // 关闭所有客户端连接
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (clients[i].is_connected) {
            closesocket(clients[i].socket);
        }
    }
    
    // 关闭服务器套接字
    closesocket(server_socket);
    
    // 清理Winsock
    WSACleanup();
}

int main(int argc, char* argv[]) {
    int port = 8888; // 默认端口
    
    // 检查命令行参数
    if (argc > 1) {
        port = atoi(argv[1]);
    }
    
    // 初始化
    if (!init_winsock()) {
        return 1;
    }
    
    if (!create_server_socket(port)) {
        WSACleanup();
        return 1;
    }
    
    init_clients();
    
    printf("聊天服务器已启动，按Ctrl+C退出...\n");
    
    // 运行服务器
    run_server();
    
    // 清理
    cleanup();
    
    return 0;
}